// Copyright 2016 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/metrics/persistent_sample_map.h"

#include "base/logging.h"
#include "base/memory/ptr_util.h"
#include "base/metrics/persistent_histogram_allocator.h"
#include "base/stl_util.h"

namespace base {

typedef HistogramBase::Count Count;
typedef HistogramBase::Sample Sample;

namespace {

    // An iterator for going through a PersistentSampleMap. The logic here is
    // identical to that of SampleMapIterator but with different data structures.
    // Changes here likely need to be duplicated there.
    class PersistentSampleMapIterator : public SampleCountIterator {
    public:
        typedef std::map<HistogramBase::Sample, HistogramBase::Count*>
            SampleToCountMap;

        explicit PersistentSampleMapIterator(const SampleToCountMap& sample_counts);
        ~PersistentSampleMapIterator() override;

        // SampleCountIterator:
        bool Done() const override;
        void Next() override;
        void Get(HistogramBase::Sample* min,
            HistogramBase::Sample* max,
            HistogramBase::Count* count) const override;

    private:
        void SkipEmptyBuckets();

        SampleToCountMap::const_iterator iter_;
        const SampleToCountMap::const_iterator end_;
    };

    PersistentSampleMapIterator::PersistentSampleMapIterator(
        const SampleToCountMap& sample_counts)
        : iter_(sample_counts.begin())
        , end_(sample_counts.end())
    {
        SkipEmptyBuckets();
    }

    PersistentSampleMapIterator::~PersistentSampleMapIterator() { }

    bool PersistentSampleMapIterator::Done() const
    {
        return iter_ == end_;
    }

    void PersistentSampleMapIterator::Next()
    {
        DCHECK(!Done());
        ++iter_;
        SkipEmptyBuckets();
    }

    void PersistentSampleMapIterator::Get(Sample* min,
        Sample* max,
        Count* count) const
    {
        DCHECK(!Done());
        if (min)
            *min = iter_->first;
        if (max)
            *max = iter_->first + 1;
        if (count)
            *count = *iter_->second;
    }

    void PersistentSampleMapIterator::SkipEmptyBuckets()
    {
        while (!Done() && *iter_->second == 0) {
            ++iter_;
        }
    }

    // This structure holds an entry for a PersistentSampleMap within a persistent
    // memory allocator. The "id" must be unique across all maps held by an
    // allocator or they will get attached to the wrong sample map.
    struct SampleRecord {
        uint64_t id; // Unique identifier of owner.
        Sample value; // The value for which this record holds a count.
        Count count; // The count associated with the above value.
    };

    // The type-id used to identify sample records inside an allocator.
    const uint32_t kTypeIdSampleRecord = 0x8FE6A69F + 1; // SHA1(SampleRecord) v1

} // namespace

PersistentSampleMap::PersistentSampleMap(
    uint64_t id,
    PersistentHistogramAllocator* allocator,
    Metadata* meta)
    : HistogramSamples(id, meta)
    , allocator_(allocator)
{
}

PersistentSampleMap::~PersistentSampleMap()
{
    if (records_)
        records_->Release(this);
}

void PersistentSampleMap::Accumulate(Sample value, Count count)
{
    *GetOrCreateSampleCountStorage(value) += count;
    IncreaseSum(static_cast<int64_t>(count) * value);
    IncreaseRedundantCount(count);
}

Count PersistentSampleMap::GetCount(Sample value) const
{
    // Have to override "const" to make sure all samples have been loaded before
    // being able to know what value to return.
    Count* count_pointer = const_cast<PersistentSampleMap*>(this)->GetSampleCountStorage(value);
    return count_pointer ? *count_pointer : 0;
}

Count PersistentSampleMap::TotalCount() const
{
    // Have to override "const" in order to make sure all samples have been
    // loaded before trying to iterate over the map.
    const_cast<PersistentSampleMap*>(this)->ImportSamples(-1, true);

    Count count = 0;
    for (const auto& entry : sample_counts_) {
        count += *entry.second;
    }
    return count;
}

std::unique_ptr<SampleCountIterator> PersistentSampleMap::Iterator() const
{
    // Have to override "const" in order to make sure all samples have been
    // loaded before trying to iterate over the map.
    const_cast<PersistentSampleMap*>(this)->ImportSamples(-1, true);
    return WrapUnique(new PersistentSampleMapIterator(sample_counts_));
}

// static
PersistentMemoryAllocator::Reference
PersistentSampleMap::GetNextPersistentRecord(
    PersistentMemoryAllocator::Iterator& iterator,
    uint64_t* sample_map_id)
{
    PersistentMemoryAllocator::Reference ref = iterator.GetNextOfType(kTypeIdSampleRecord);
    const SampleRecord* record = iterator.GetAsObject<SampleRecord>(ref, kTypeIdSampleRecord);
    if (!record)
        return 0;

    *sample_map_id = record->id;
    return ref;
}

// static
PersistentMemoryAllocator::Reference
PersistentSampleMap::CreatePersistentRecord(
    PersistentMemoryAllocator* allocator,
    uint64_t sample_map_id,
    Sample value)
{
    PersistentMemoryAllocator::Reference ref = allocator->Allocate(sizeof(SampleRecord), kTypeIdSampleRecord);
    SampleRecord* record = allocator->GetAsObject<SampleRecord>(ref, kTypeIdSampleRecord);

    if (!record) {
        NOTREACHED() << "full=" << allocator->IsFull()
                     << ", corrupt=" << allocator->IsCorrupt();
        return 0;
    }

    record->id = sample_map_id;
    record->value = value;
    record->count = 0;
    allocator->MakeIterable(ref);
    return ref;
}

bool PersistentSampleMap::AddSubtractImpl(SampleCountIterator* iter,
    Operator op)
{
    Sample min;
    Sample max;
    Count count;
    for (; !iter->Done(); iter->Next()) {
        iter->Get(&min, &max, &count);
        if (min + 1 != max)
            return false; // SparseHistogram only supports bucket with size 1.

        *GetOrCreateSampleCountStorage(min) += (op == HistogramSamples::ADD) ? count : -count;
    }
    return true;
}

Count* PersistentSampleMap::GetSampleCountStorage(Sample value)
{
    // If |value| is already in the map, just return that.
    auto it = sample_counts_.find(value);
    if (it != sample_counts_.end())
        return it->second;

    // Import any new samples from persistent memory looking for the value.
    return ImportSamples(value, false);
}

Count* PersistentSampleMap::GetOrCreateSampleCountStorage(Sample value)
{
    // Get any existing count storage.
    Count* count_pointer = GetSampleCountStorage(value);
    if (count_pointer)
        return count_pointer;

    // Create a new record in persistent memory for the value. |records_| will
    // have been initialized by the GetSampleCountStorage() call above.
    DCHECK(records_);
    PersistentMemoryAllocator::Reference ref = records_->CreateNew(value);
    if (!ref) {
        // If a new record could not be created then the underlying allocator is
        // full or corrupt. Instead, allocate the counter from the heap. This
        // sample will not be persistent, will not be shared, and will leak...
        // but it's better than crashing.
        count_pointer = new Count(0);
        sample_counts_[value] = count_pointer;
        return count_pointer;
    }

    // A race condition between two independent processes (i.e. two independent
    // histogram objects sharing the same sample data) could cause two of the
    // above records to be created. The allocator, however, forces a strict
    // ordering on iterable objects so use the import method to actually add the
    // just-created record. This ensures that all PersistentSampleMap objects
    // will always use the same record, whichever was first made iterable.
    // Thread-safety within a process where multiple threads use the same
    // histogram object is delegated to the controlling histogram object which,
    // for sparse histograms, is a lock object.
    count_pointer = ImportSamples(value, false);
    DCHECK(count_pointer);
    return count_pointer;
}

PersistentSampleMapRecords* PersistentSampleMap::GetRecords()
{
    // The |records_| pointer is lazily fetched from the |allocator_| only on
    // first use. Sometimes duplicate histograms are created by race conditions
    // and if both were to grab the records object, there would be a conflict.
    // Use of a histogram, and thus a call to this method, won't occur until
    // after the histogram has been de-dup'd.
    if (!records_)
        records_ = allocator_->UseSampleMapRecords(id(), this);
    return records_;
}

Count* PersistentSampleMap::ImportSamples(Sample until_value,
    bool import_everything)
{
    Count* found_count = nullptr;
    PersistentMemoryAllocator::Reference ref;
    PersistentSampleMapRecords* records = GetRecords();
    while ((ref = records->GetNext()) != 0) {
        SampleRecord* record = records->GetAsObject<SampleRecord>(ref, kTypeIdSampleRecord);
        if (!record)
            continue;

        DCHECK_EQ(id(), record->id);

        // Check if the record's value is already known.
        if (!ContainsKey(sample_counts_, record->value)) {
            // No: Add it to map of known values.
            sample_counts_[record->value] = &record->count;
        } else {
            // Yes: Ignore it; it's a duplicate caused by a race condition -- see
            // code & comment in GetOrCreateSampleCountStorage() for details.
            // Check that nothing ever operated on the duplicate record.
            DCHECK_EQ(0, record->count);
        }

        // Check if it's the value being searched for and, if so, keep a pointer
        // to return later. Stop here unless everything is being imported.
        // Because race conditions can cause multiple records for a single value,
        // be sure to return the first one found.
        if (record->value == until_value) {
            if (!found_count)
                found_count = &record->count;
            if (!import_everything)
                break;
        }
    }

    return found_count;
}

} // namespace base
